import streamlit as st
import cv2
import numpy as np
import torch
from PIL import Image
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from streamlit_drawable_canvas import st_canvas

# 页面配置
st.set_page_config(
    page_title="何明钊的手写数字识别系统",
    page_icon="✍️",
    layout="wide"
)

# 自定义CSS样式
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .title {
        color: #2c3e50;
        text-align: center;
        padding: 1.5rem;
        border-bottom: 2px solid #eee;
        margin-bottom: 2rem;
    }
    .stButton>button {
        width: 100%;
        border-radius: 5px;
        height: 3em;
        background-color: #2c3e50;
        color: white;
    }
    .prediction-result {
        padding: 1rem;
        border-radius: 5px;
        background-color: #f8f9fa;
        margin: 1rem 0;
    }
    .confidence {
        color: #2c3e50;
        font-size: 0.9em;
    }
    .digit-display {
        text-align: center;
        padding: 0.5rem;
        background-color: #f8f9fa;
        border-radius: 5px;
        margin: 0.5rem;
    }
    .canvas-container {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
    .help-text {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 5px;
        margin-top: 1rem;
    }
    .stSelectbox {
        margin-bottom: 1rem;
    }
    </style>
    """, unsafe_allow_html=True)


# 模型定义
class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x


# 加载模型
@st.cache_resource
def load_model():
    model = CNNModel()
    model.load_state_dict(torch.load('E:/mnist/mnist_model.pth', map_location=torch.device('cpu')))
    model.eval()
    return model


# 图像预处理
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])


def process_canvas_image(image_data):
    # 转换为灰度图
    gray = cv2.cvtColor(image_data, cv2.COLOR_RGBA2GRAY)
    # 二值化，注意这里需要反转颜色（因为画布是白字黑底）
    _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
    # 向左旋转90度
    binary = np.rot90(binary)
    # 再向左旋转90度回来
    binary = np.rot90(binary, k=3)  # k=3相当于向右旋转90度
    return binary


def split_digits(image):
    img_size = 28
    kernel = np.ones((3, 3), np.uint8)

    # 确保图像是灰度图
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化处理
    _, thresh_img = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY_INV)

    # 膨胀处理
    dilated = cv2.dilate(thresh_img, kernel, iterations=1)

    # 查找轮廓
    contours, _ = cv2.findContours(dilated.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 按x坐标排序轮廓
    contours = sorted(contours, key=lambda x: cv2.boundingRect(x)[0])

    digit_images = []
    for contour in contours:
        [x, y, w, h] = cv2.boundingRect(contour)

        # 过滤掉太小的轮廓
        if w * h < 100:
            continue

        # 提取数字区域
        digit = thresh_img[y:y + h, x:x + w]

        # 添加边框
        pad_h = int(h * 0.35)
        pad_w = int(w * 0.35)
        digit = cv2.copyMakeBorder(digit, pad_h, pad_h, pad_w, pad_w,
                                   cv2.BORDER_CONSTANT, value=0)

        # 调整大小
        digit = cv2.resize(digit, (img_size, img_size))

        digit_images.append(digit)

    return digit_images


# 预测函数
def predict_digit(image, model):
    with torch.no_grad():
        image_tensor = transform(Image.fromarray(image)).unsqueeze(0)
        output = model(image_tensor)
        probabilities = F.softmax(output, dim=1)
        pred_prob, predicted = torch.max(probabilities, 1)
        return predicted.item(), pred_prob.item() * 100


def main():
    st.markdown("<h1 class='title'>✍️ 何明钊的手写数字识别系统</h1>", unsafe_allow_html=True)

    # 加载模型
    model = load_model()

    # 使用container来组织布局
    with st.container():
        col1, col2, col3 = st.columns([1, 3, 1])

        with col1:
            st.markdown("### 🛠️ 设置")
            mode = st.selectbox(
                "选择模式",
                ["单个数字识别", "数字串识别"]
            )
            input_method = st.selectbox(
                "选择输入方式",
                ["手写输入", "上传图片"]
            )

        with col2:
            if mode == "单个数字识别":
                st.markdown("<h3 style='text-align: center;'>✏️ 单个数字识别</h3>", unsafe_allow_html=True)

                if input_method == "手写输入":
                    # 添加画笔粗细选择器
                    brush_size = st.slider("选择画笔粗细", 1, 50, 20)

                    st.markdown("<div class='canvas-container'>", unsafe_allow_html=True)
                    canvas_result = st_canvas(
                        stroke_width=brush_size,  # 使用选择的画笔粗细
                        stroke_color="#fff",
                        background_color="#000",
                        height=280,
                        width=280,
                        drawing_mode="freedraw",
                        key="canvas",
                    )
                    st.markdown("</div>", unsafe_allow_html=True)

                    if canvas_result.image_data is not None:
                        image = canvas_result.image_data
                        if image is not None:
                            gray = cv2.cvtColor(image, cv2.COLOR_RGBA2GRAY)
                            digit, confidence = predict_digit(gray, model)

                            st.markdown("<div class='prediction-result'>", unsafe_allow_html=True)
                            col_a, col_b = st.columns(2)
                            with col_a:
                                st.markdown(f"<h2 style='text-align: center;'>预测结果: {digit}</h2>",
                                            unsafe_allow_html=True)
                            with col_b:
                                st.markdown(f"<p class='confidence'>置信度: {confidence:.2f}%</p>",
                                            unsafe_allow_html=True)
                            st.markdown("</div>", unsafe_allow_html=True)

                else:
                    uploaded_file = st.file_uploader("上传单个数字图片", type=['png', 'jpg', 'jpeg'])
                    if uploaded_file is not None:
                        image = Image.open(uploaded_file).convert('L')
                        image = np.array(image)
                        st.image(image, caption='上传的图片', width=200)

                        digit, confidence = predict_digit(image, model)
                        st.markdown("<div class='prediction-result'>", unsafe_allow_html=True)
                        st.markdown(f"<h2 style='text-align: center;'>预测结果: {digit}</h2>",
                                    unsafe_allow_html=True)
                        st.markdown(f"<p class='confidence'>置信度: {confidence:.2f}%</p>",
                                    unsafe_allow_html=True)
                        st.markdown("</div>", unsafe_allow_html=True)

            else:
                st.markdown("<h3 style='text-align: center;'>📝 数字串识别</h3>", unsafe_allow_html=True)

                if input_method == "手写输入":
                    brush_size = st.slider("选择画笔粗细", 1, 50, 20)

                    st.markdown("<div class='canvas-container'>", unsafe_allow_html=True)
                    canvas_result = st_canvas(
                        stroke_width=brush_size,
                        stroke_color="#fff",
                        background_color="#000",
                        height=200,
                        width=600,
                        drawing_mode="freedraw",
                        key="canvas_string",
                    )
                    st.markdown("</div>", unsafe_allow_html=True)

                    if canvas_result.image_data is not None:
                        image = canvas_result.image_data
                        if image is not None:
                            processed_image = process_canvas_image(image)
                            digits = split_digits(processed_image)

                            if digits:
                                st.markdown("<div class='prediction-result'>", unsafe_allow_html=True)
                                st.markdown("<h4>分割结果：</h4>", unsafe_allow_html=True)
                                cols = st.columns(len(digits))
                                for idx, digit in enumerate(digits):
                                    with cols[idx]:
                                        st.markdown(f"<div class='digit-display'>", unsafe_allow_html=True)
                                        st.image(digit, width=50)
                                        st.markdown(f"数字 {idx + 1}", unsafe_allow_html=True)
                                        st.markdown("</div>", unsafe_allow_html=True)

                                results = []
                                result_str = ""
                                for digit_image in digits:
                                    digit, confidence = predict_digit(digit_image, model)
                                    results.append((digit, confidence))
                                    result_str += str(digit)

                                st.markdown("<h4>识别结果：</h4>", unsafe_allow_html=True)
                                st.markdown(f"<h2 style='text-align: center;'>{result_str}</h2>",
                                            unsafe_allow_html=True)

                                for i, (digit, conf) in enumerate(results):
                                    st.markdown(f"<p class='confidence'>第{i + 1}个数字: {digit}, "
                                                f"置信度: {conf:.2f}%</p>", unsafe_allow_html=True)
                                st.markdown("</div>", unsafe_allow_html=True)
                            else:
                                st.error("未能检测到数字，请重新书写")

                else:
                    uploaded_file = st.file_uploader("上传数字串图片", type=['png', 'jpg', 'jpeg'])
                    if uploaded_file is not None:
                        image = Image.open(uploaded_file).convert('L')
                        image = np.array(image)
                        st.image(image, caption='上传的图片', width=400)

                        digits = split_digits(image)
                        if digits:
                            st.markdown("<div class='prediction-result'>", unsafe_allow_html=True)
                            st.markdown("<h4>分割结果：</h4>", unsafe_allow_html=True)
                            cols = st.columns(len(digits))
                            for idx, digit in enumerate(digits):
                                with cols[idx]:
                                    st.markdown(f"<div class='digit-display'>", unsafe_allow_html=True)
                                    st.image(digit, width=50)
                                    st.markdown(f"数字 {idx + 1}", unsafe_allow_html=True)
                                    st.markdown("</div>", unsafe_allow_html=True)

                            results = []
                            result_str = ""
                            for digit_image in digits:
                                digit, confidence = predict_digit(digit_image, model)
                                results.append((digit, confidence))
                                result_str += str(digit)

                            st.markdown("<h4>识别结果：</h4>", unsafe_allow_html=True)
                            st.markdown(f"<h2 style='text-align: center;'>{result_str}</h2>",
                                        unsafe_allow_html=True)

                            for i, (digit, conf) in enumerate(results):
                                st.markdown(f"<p class='confidence'>第{i + 1}个数字: {digit}, "
                                            f"置信度: {conf:.2f}%</p>", unsafe_allow_html=True)
                            st.markdown("</div>", unsafe_allow_html=True)
                        else:
                            st.error("未能检测到数字，请重新上传")

        with col3:
            st.markdown("### ❓ 帮助")
            st.markdown("""
            <div class='help-text'>
            <b>使用说明：</b><br>
            1. 选择识别模式<br>
            2. 选择输入方式<br>
            3. 在画布上书写或上传图片<br>
            4. 查看识别结果<br><br>
            <b>注意事项：</b><br>
            1. 书写时请保持字迹清晰<br>
            2. 数字串识别时请保持适当间距<br>
            3. 尽量保持数字大小一致<br>
            4. 上传图片请确保图片清晰度
            </div>
            """, unsafe_allow_html=True)

            # 添加版本信息
            st.markdown("""
            <div style='position: fixed; bottom: 20px; right: 20px; text-align: right;'>
            <p style='color: #666; font-size: 0.8em;'>
            Version 1.0<br>
            Powered by Streamlit
            </p>
            </div>
            """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()